﻿# -*- coding: utf-8 -*-
"""
@author: Junxiao Song  Haoran Wang
"""

import numpy as np
import copy


class Board(object):
    """board for the game"""
    """
    如果棋盘是
    -1  0 -1  0
     1  0  1  1
     1  1  1  0
     1  0 -1  0
    则对应的states是
    {0: 1, 1: 0, 2: -1, 3: 0, 4: 1, 5: 1, 6: 1, 7: 0, 8: 1, 9: 0, 10: 1, 11: 1, 12: -1, 13: 0, 14: -1, 15: 0}
    """

    def __init__(self, **kwargs):
        self.width = int(kwargs.get('width', 4))
        self.height = int(kwargs.get('height', 4))
        # board states stored as a dict,
        # key: move as location on the board,
        # value: player as pieces type
        self.states = {}
        self.history = [];
        # need how many pieces in a row to win
        self.n_in_row = int(kwargs.get('n_in_row', 4))
        self.players = [1, -1]  # player1 and player2
        self.cases = ["aabo","baao","abao","babo","abbb","aabb","abba","baab"]

    def init_board(self, start_player=0):
        if self.width < self.n_in_row or self.height < self.n_in_row:
            raise Exception('board width and height can not be '
                            'less than {}'.format(self.n_in_row))
        self.current_player = self.players[start_player]  # start player
        self.states = {0: 0, 1: 0, 2: 0, 3: 0, 4: 0, 5: 0, 6: 0, 7: 0, 8: 0, 9: 0, 10: 0, 11: 0, 12: 0, 13: 0, 14: 0, 15: 0}
        # self.states = {0: -1, 1: 0, 2: -1, 3: 0, 4: -1, 5: -1, 6: 1, 7: -1, 8: 1, 9: 1, 10: 1, 11: -1, 12: 0, 13: 1, 14: 0, 15: 0}
        # keep available moves in a list
        self.availables = []
        for key,value in self.states.items():
            if value == 0:
                self.availables.append(key)
        self.last_move = -1
        self.history.append(str(self.states))

    def move_to_location(self, move):
        """
        4*4 board's moves like:
        8 9 ......
        4 5 6 7
        0 1 2 3
        and move 7's location is (2,3)
        """
        h = move // self.width #落子所在的高度
        w = move % self.width #落子所在的宽度
        return [h, w]

    def location_to_move(self, location):
        if len(location) != 2:
            return -1
        h = location[0]
        w = location[1]
        move = h * self.width + w
        if move not in range(self.width * self.height):
            return -1
        return move

    def current_state(self):
        """return the board state from the perspective of the current player.
        state shape: 4*width*height
        """

        square_state = np.zeros((4, self.width, self.height))
        if self.states:
            moves, players = np.array(list(zip(*self.states.items())))
            move_curr = moves[players == self.current_player]
            move_oppo = moves[players != self.current_player]
            square_state[0][move_curr // self.width,
                            move_curr % self.height] = 1.0
            square_state[1][move_oppo // self.width,
                            move_oppo % self.height] = 1.0
            # indicate the last move location
            square_state[2][self.last_move // self.width,
                            self.last_move % self.height] = 1.0
        if len(self.states) % 2 == 0:
            square_state[3][:, :] = 1.0  # indicate the colour to play
        return square_state[:, ::-1, :]
    
    #提取特征字符串
    @staticmethod
    def get_character_string(vector,holder):
        """
            Input:
                vector: 要转换为字符串的向量
                holder: 向量中的执手子
            Returns:
                cs: 向量的特征字符串
                    例如向量[1,1,-1,0]，执手子为1
                    特征字符串就是 aabo
        """
        cs = ""
        for i in vector:
            if i == holder:
                cs += 'a'
            elif i == 0:
                cs += 'o'
            else:
                cs += 'b'
        return cs
    
    def get_row(self,move):
        """
            Input: 
                move: 落子的位置
            Return:
                从states中读取到的落子行
        """
        row = []
        states = self.states
        width = self.width
        w = move // width
        for state in states:
            if state // width == w:
                row.append(states[state])
        return np.array(row)
                
    def set_row(self,row,move):
        """
            Input: 
                row:  消字后的行
                move: 上次落子的位置
        """
        states = self.states
        width = self.width
        w = move // width
        i = 0
        for state in states:
            if state // width == w and i < len(row):
                states[state] = row[i]
                i += 1
                    

    def get_col(self,move):
        """
            Input: 
                move: 落子的位置
            Return:
                从sates中读取到的落子列
        """
        col = []
        states = self.states
        height = self.height
        h = move % height
        for state in states:
            if state % height == h:
                col.append(states[state])
        return np.array(col)
    
    def set_col(self,col,move):
        """
            Input: 
                col:  消字后的列
                move: 上次落子的位置
        """
        states = self.states
        height = self.height
        h = move % height
        j = 0
        for state in states:
            if state % height == h and j < len(col):
                states[state] = col[j]
                j += 1
    
    def eliminate(self,move):
        """
            消子的方法，参考特征字符串cases进行消字
            Input: 落子所在的位置
        """
        cases = self.cases
        states = self.states
        
        row = self.get_row(move)
        col = self.get_col(move)
        
        holder = states[move]
        
        cs_row = self.get_character_string(row,holder)
        cs_col = self.get_character_string(col,holder)

        if cs_row in self.cases or cs_row[::-1] in self.cases:
            row = np.where(row == -holder,0,row)
            self.set_row(row,move)
        if cs_col in self.cases or cs_col[::-1] in self.cases:
            col = np.where(col == -holder,0,col)
            self.set_col(col,move)
        return self;
        
    
    def do_move(self, move):
        if move in self.availables:
            self.states[move] = self.current_player
            states = self.states
            self = self.eliminate(move)
            self.availables.remove(move)
            #更新availables
            for state in states:
                if states[state] == 0:
                    if not state in self.availables:
                        newBoard = copy.deepcopy(self)
                        newBoard.states[state] = self.current_player * -1
                        newBoard.eliminate(state)
                        newStateStr = str(newBoard.states)
                        # print(newStateStr)
                        if not newStateStr in self.history:
                            # 禁止全局同型循环
                            self.availables.append(state)
            for state in self.availables:
                newBoard = copy.deepcopy(self)
                newBoard.states[state] = self.current_player * -1
                newBoard.eliminate(state)
                newStateStr = str(newBoard.states)
                # print(newStateStr)
                if newStateStr in self.history:
                    # 禁止全局同型循环
                    self.availables.remove(state)
            self.current_player = (
                self.players[0] if self.current_player == self.players[1]
                else self.players[1]
            )
            self.last_move = move
            self.history.append(str(self.states))

    def has_a_winner(self):
        width = self.width
        height = self.height
        states = self.states
        n = self.n_in_row

        # 已经落子的位置
        moved = [key for key, value in states.items() if value != 0]
        #棋子总数小于4个，无法构成一条驴
        if len(moved) < self.n_in_row : 
            return False, -1
        
        panel = np.array([[0]*width]*height)
        
        
        
        white_win = "aaaa"
        black_win = "bbbb"
        
        for m in moved:
            h = m // width
            w = m % width
            panel[h][w] = states[m]
          
        css = [] # character string
        for row in panel:
            css.append(Board.get_character_string(row, 1))
        for col in panel.T:
            css.append(Board.get_character_string(col, 1))

        if white_win in css: #白字一条驴
            #print("白子一条吕")
            return True, 1
        if black_win in css: #黑子一条驴
            #print("黑子一条吕")
            return True,-1
        if len(moved) == width * height : #棋盘下满
            if panel.sum() > 0: #白子多，白子赢
                #print("白子多，白子赢")
                return True, 1
            if panel.sum() < 0: #黑子多，黑子赢
                #print("黑子多，黑子赢")
                return True,-1
        return False, -1

    def game_end(self):
        """Check whether the game is ended or not"""
        win, winner = self.has_a_winner()
        if win: #某玩家胜利
            return True, winner
        elif not len(self.availables): #平局
            return True, 0
        return False, 0 #游戏未结束

    def get_current_player(self):
        return self.current_player


class Game(object):
    """game server"""

    def __init__(self, board, **kwargs):
        
        self.board = board

    def graphic(self, board, player1, player2):
        """Draw the board and show game info"""
        width = board.width
        height = board.height
        
        print("Player", player1, "with X".rjust(3))
        print("Player", player2, "with O".rjust(3))
        print()
        for x in range(width):
            print("{0:8}".format(x), end='')
        print('\r\n')
        for i in range(height - 1, -1, -1):
            print("{0:4d}".format(i), end='')
            for j in range(width):
                loc = i * width + j
                p = board.states.get(loc, 0)
                if p == player1:
                    print('X'.center(8), end='')
                elif p == player2:
                    print('O'.center(8), end='')
                else:
                    print('_'.center(8), end='')
            print('\r\n\r\n')

    def start_play(self, player1, player2, start_player=0, is_shown=1):
        """start a game between two players"""
        if start_player not in (0, 1):
            raise Exception('start_player should be either 0 (player1 first) '
                            'or 1 (player2 first)')
        self.board.init_board(start_player)
        p1, p2 = self.board.players
        player1.set_player_ind(p1)
        player2.set_player_ind(p2)
        players = {p1: player1, p2: player2}
        if is_shown:
            self.graphic(self.board, player1.player, player2.player)
        while True:
            current_player = self.board.get_current_player()
            player_in_turn = players[current_player]
            move = player_in_turn.get_action(self.board)
            self.board.do_move(move)
            if is_shown:
                self.graphic(self.board, player1.player, player2.player)
            end, winner = self.board.game_end()
            if end:
                if is_shown:
                    if winner != 0:
                        print("Game end. Winner is", players[winner])
                    else:
                        print("Game end. Tie")
                return winner

    def start_self_play(self, player, is_shown=0, temp=1e-3):
        """ start a self-play game using a MCTS player, reuse the search tree,
        and store the self-play data: (state, mcts_probs, z) for training
        """
        self.board.init_board()
        p1, p2 = self.board.players
        states, mcts_probs, current_players = [], [], []
        while True:
            move, move_probs = player.get_action(self.board,
                                                 temp=temp,
                                                 return_prob=1)
            # store the data
            states.append(self.board.current_state())
            mcts_probs.append(move_probs)
            current_players.append(self.board.current_player)
            # perform a move
            self.board.do_move(move)
            if is_shown:
                self.graphic(self.board, p1, p2)
            end, winner = self.board.game_end()
            if end:
                # winner from the perspective of the current player of each state
                winners_z = np.zeros(len(current_players))
                if winner != -1:
                    winners_z[np.array(current_players) == winner] = 1.0
                    winners_z[np.array(current_players) != winner] = -1.0
                # reset MCTS root node
                player.reset_player()
                if is_shown:
                    if winner != -1:
                        print("Game end. Winner is player:", winner)
                    else:
                        print("Game end. Tie")
                return winner, zip(states, mcts_probs, winners_z)
